import pygame
import random
import sys

# 初始化
pygame.init()

# 屏幕大小
WIDTH, HEIGHT = 600, 400
CELL_SIZE = 20
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("贪吃蛇")

# 颜色
BLACK = (0, 0, 0)
GREEN = (0, 200, 0)
RED = (200, 0, 0)
WHITE = (255, 255, 255)

# 时钟
clock = pygame.time.Clock()

# 初始蛇
snake = [(100, 100), (80, 100), (60, 100)]
direction = "RIGHT"

# 生成食物
def random_food():
    x = random.randint(0, (WIDTH - CELL_SIZE) // CELL_SIZE) * CELL_SIZE
    y = random.randint(0, (HEIGHT - CELL_SIZE) // CELL_SIZE) * CELL_SIZE
    return (x, y)

food = random_food()

# 游戏主循环
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP and direction != "DOWN":
                direction = "UP"
            elif event.key == pygame.K_DOWN and direction != "UP":
                direction = "DOWN"
            elif event.key == pygame.K_LEFT and direction != "RIGHT":
                direction = "LEFT"
            elif event.key == pygame.K_RIGHT and direction != "LEFT":
                direction = "RIGHT"

    # 移动蛇
    x, y = snake[0]
    if direction == "UP":
        y -= CELL_SIZE
    elif direction == "DOWN":
        y += CELL_SIZE
    elif direction == "LEFT":
        x -= CELL_SIZE
    elif direction == "RIGHT":
        x += CELL_SIZE

    new_head = (x, y)

    # 判断碰撞
    if (
        x < 0 or x >= WIDTH or
        y < 0 or y >= HEIGHT or
        new_head in snake
    ):
        print("游戏结束! 得分:", len(snake) - 3)
        pygame.quit()
        sys.exit()

    # 吃食物
    snake.insert(0, new_head)
    if new_head == food:
        food = random_food()
    else:
        snake.pop()

    # 画面更新
    screen.fill(BLACK)
    for s in snake:
        pygame.draw.rect(screen, GREEN, (s[0], s[1], CELL_SIZE, CELL_SIZE))
    pygame.draw.rect(screen, RED, (food[0], food[1], CELL_SIZE, CELL_SIZE))

    # 显示分数
    font = pygame.font.SysFont(None, 30)
    score = font.render(f"得分: {len(snake)-3}", True, WHITE)
    screen.blit(score, (10, 10))

    pygame.display.flip()
    clock.tick(10)  # 控制速度
